# -*- coding: utf-8 -*-
"""
MACD成交量放大强势筛选器 - 成交量放大且其他指标配合
筛选规则：
1. macd_dif > macd_dea（MACD金叉或强势）
2. volume >= 1.5 * avg_vol_20（成交量明显放大）
3. rsi > 50（相对强势）
4. close > ma20 且 ma20 > ma60（短期均线向上且处于多头排列）

策略说明：
- 成交量显著放大，显示资金活跃介入
- MACD处于强势状态，趋势向上
- RSI超过50，价格相对强势
- 均线多头排列，短期趋势明确向上
- 综合多个指标确认强势上涨信号
"""
import logging
import pandas as pd
import numpy as np
from datetime import datetime
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from database import DatabaseManager
from config import DATABASE_CONFIG
import warnings

# 忽略pandas警告
warnings.filterwarnings('ignore', category=pd.errors.SettingWithCopyWarning)
warnings.filterwarnings('ignore', category=FutureWarning)

# 配置日志
# 获取项目根目录
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
LOG_DIR = os.path.join(PROJECT_ROOT, 'logs')
os.makedirs(LOG_DIR, exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler(os.path.join(LOG_DIR, 'macd_volume_strength_screener.log'), encoding='utf-8')
    ]
)
logger = logging.getLogger(__name__)


class MACDVolumeStrengthScreener:
    """MACD成交量放大强势筛选器"""
    
    def __init__(self):
        self.db = DatabaseManager(DATABASE_CONFIG)
    
    def get_stock_daily_data(self, stock_code: str, days: int = 200) -> pd.DataFrame:
        """获取股票日线数据"""
        sql = """
        SELECT trade_date, open_price, high_price, low_price, close_price, volume, change_pct
        FROM new_stock_daily 
        WHERE stock_code = %s 
        ORDER BY trade_date DESC 
        LIMIT %s
        """
        try:
            results = self.db.execute_query(sql, (stock_code, days))
            if not results:
                return pd.DataFrame()
            
            df = pd.DataFrame(results)
            df['trade_date'] = pd.to_datetime(df['trade_date'])
            
            # 转换为float类型
            price_columns = ['open_price', 'high_price', 'low_price', 'close_price']
            for col in price_columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')
            
            df['volume'] = pd.to_numeric(df['volume'], errors='coerce')
            df['change_pct'] = pd.to_numeric(df['change_pct'], errors='coerce')
            
            df = df.sort_values('trade_date').reset_index(drop=True)
            return df
        except Exception as e:
            logger.error(f"获取股票 {stock_code} 数据失败: {e}")
            return pd.DataFrame()
    
    def calculate_ma(self, df: pd.DataFrame, periods: list = [20, 60]) -> pd.DataFrame:
        """计算移动平均线"""
        if df.empty:
            return df
        
        for period in periods:
            if len(df) >= period:
                df[f'ma{period}'] = df['close_price'].rolling(window=period).mean()
        
        return df
    
    def calculate_ema(self, df: pd.DataFrame, period: int) -> pd.Series:
        """计算指数移动平均线"""
        return df['close_price'].ewm(span=period, adjust=False).mean()
    
    def calculate_macd(self, df: pd.DataFrame, fast_period: int = 12, 
                      slow_period: int = 26, signal_period: int = 9) -> pd.DataFrame:
        """
        计算MACD指标
        DIF = EMA(12) - EMA(26)
        DEA = EMA(DIF, 9)
        HIST = (DIF - DEA) * 2
        """
        if df.empty or len(df) < slow_period:
            return df
        
        # 计算EMA
        ema_fast = self.calculate_ema(df, fast_period)
        ema_slow = self.calculate_ema(df, slow_period)
        
        # 计算DIF
        df['macd_dif'] = ema_fast - ema_slow
        
        # 计算DEA
        df['macd_dea'] = df['macd_dif'].ewm(span=signal_period, adjust=False).mean()
        
        # 计算HIST（柱状图）
        df['macd_hist'] = (df['macd_dif'] - df['macd_dea']) * 2
        
        return df
    
    def calculate_rsi(self, df: pd.DataFrame, period: int = 14) -> pd.DataFrame:
        """计算RSI相对强弱指数"""
        if df.empty or len(df) < period:
            return df
        
        # 计算价格变化
        delta = df['close_price'].diff()
        
        # 分离上涨和下跌
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        
        # 计算平均涨幅和跌幅
        avg_gain = gain.rolling(window=period).mean()
        avg_loss = loss.rolling(window=period).mean()
        
        # 避免除零
        avg_loss = avg_loss.replace(0, 0.0001)
        
        # 计算RS和RSI
        rs = avg_gain / avg_loss
        df['rsi'] = 100 - (100 / (1 + rs))
        
        return df
    
    def calculate_indicators(self, df: pd.DataFrame) -> pd.DataFrame:
        """计算所有技术指标"""
        if df.empty or len(df) < 70:
            return df
        
        # 计算移动平均线
        df = self.calculate_ma(df, periods=[20, 60])
        
        # 计算MACD
        df = self.calculate_macd(df)
        
        # 计算RSI
        df = self.calculate_rsi(df)
        
        # 计算20日平均成交量
        df['avg_vol_20'] = df['volume'].rolling(window=20).mean()
        
        # 计算前一日的MACD DIF和DEA（用于判断金叉时间）
        df['prev_macd_dif'] = df['macd_dif'].shift(1)
        df['prev_macd_dea'] = df['macd_dea'].shift(1)
        
        # 计算前两日的MACD DIF和DEA（用于判断金叉时间）
        df['prev2_macd_dif'] = df['macd_dif'].shift(2)
        df['prev2_macd_dea'] = df['macd_dea'].shift(2)
        
        return df
    
    def check_volume_strength_condition(self, row: pd.Series) -> bool:
        """
        检查是否满足成交量放大强势条件
        1. macd_dif > macd_dea（MACD金叉或强势）
        2. volume >= 1.5 * avg_vol_20（成交量明显放大）
        3. rsi > 50（相对强势）
        4. close > ma20 且 ma20 > ma60（短期均线向上且多头排列）
        """
        try:
            # 检查必要字段是否存在
            required_fields = ['macd_dif', 'macd_dea', 'volume', 'avg_vol_20',
                             'rsi', 'close_price', 'ma20', 'ma60',
                             'prev_macd_dif', 'prev_macd_dea', 
                             'prev2_macd_dif', 'prev2_macd_dea']
            if any(pd.isna(row[field]) for field in required_fields):
                return False
            
            # 条件1：MACD金叉或强势状态
            if row['macd_dif'] <= row['macd_dea']:
                return False
            
            # 检查金叉是否是新发生的（前1日或前2日时还未金叉）
            # 这样可以捕获最新的金叉信号
            prev1_is_golden = row['prev_macd_dif'] > row['prev_macd_dea']
            prev2_is_golden = row['prev2_macd_dif'] > row['prev2_macd_dea']
            
            # 金叉必须是新发生的：
            # 情况1：昨日未金叉，今日金叉（今日刚金叉）
            # 情况2：前日未金叉，昨日金叉（昨日刚金叉）
            is_fresh_golden_cross = (not prev1_is_golden) or (not prev2_is_golden)
            
            if not is_fresh_golden_cross:
                return False
            
            # 条件2：成交量明显放大（1.5倍以上）
            if row['volume'] < 1.5 * row['avg_vol_20']:
                return False
            
            # 条件3：RSI > 50，表示相对强势
            if row['rsi'] <= 50:
                return False
            
            # 条件4：价格在MA20之上
            if row['close_price'] <= row['ma20']:
                return False
            
            # 条件5：MA20在MA60之上（多头排列）
            if row['ma20'] <= row['ma60']:
                return False
            
            return True
            
        except Exception as e:
            logger.error(f"检查成交量强势条件失败: {e}")
            return False
    
    def screen_stock(self, stock_code: str, stock_name: str) -> dict:
        """筛选单只股票，带重试机制"""
        max_retries = 2
        for attempt in range(max_retries):
            try:
                # 获取日线数据
                df = self.get_stock_daily_data(stock_code, days=200)
                if df.empty or len(df) < 70:
                    return None
                
                # 计算技术指标
                df = self.calculate_indicators(df)
                
                # 获取最新一天的数据
                latest = df.iloc[-1]
                
                # 检查是否满足条件
                if not self.check_volume_strength_condition(latest):
                    return None
                
                # 计算量比
                vol_ratio = float(latest['volume'] / latest['avg_vol_20'])
                
                # 计算价格相对MA20的位置
                price_above_ma20_pct = float((latest['close_price'] - latest['ma20']) / latest['ma20'] * 100)
                
                # 计算MA20相对MA60的位置
                ma20_above_ma60_pct = float((latest['ma20'] - latest['ma60']) / latest['ma60'] * 100)
                
                # 返回筛选结果
                result = {
                    'stock_code': stock_code,
                    'stock_name': stock_name,
                    'trade_date': latest['trade_date'].strftime('%Y-%m-%d'),
                    'close_price': round(float(latest['close_price']), 2),
                    'rsi': round(float(latest['rsi']), 2),
                    'ma20': round(float(latest['ma20']), 2),
                    'ma60': round(float(latest['ma60']), 2),
                    'macd_dif': round(float(latest['macd_dif']), 4),
                    'macd_dea': round(float(latest['macd_dea']), 4),
                    'volume': int(latest['volume']),
                    'avg_vol_20': round(float(latest['avg_vol_20']), 2),
                    'vol_ratio': round(vol_ratio, 2),
                    'price_above_ma20_pct': round(price_above_ma20_pct, 2),
                    'ma20_above_ma60_pct': round(ma20_above_ma60_pct, 2),
                    'change_pct': round(float(latest['change_pct']), 2) if not pd.isna(latest['change_pct']) else 0,
                }
                
                return result
                
            except Exception as e:
                if attempt < max_retries - 1:
                    logger.warning(f"筛选股票 {stock_code} 失败，重试中 (尝试 {attempt + 1}/{max_retries}): {e}")
                    import time
                    time.sleep(0.5)  # 短暂延迟后重试
                else:
                    logger.error(f"筛选股票 {stock_code} 失败: {e}")
                    return None
    
    def get_all_active_stocks(self):
        """获取所有活跃股票"""
        sql = """
        SELECT stock_code, stock_name 
        FROM new_stock_info 
        WHERE status = 'active'
        ORDER BY stock_code
        """
        return self.db.execute_query(sql)
    
    def run_screening(self, limit: int = None):
        """运行筛选"""
        logger.info("=" * 80)
        logger.info("开始MACD成交量放大强势筛选 - 成交量放大且其他指标配合")
        logger.info("=" * 80)
        
        # 获取所有活跃股票
        stocks = self.get_all_active_stocks()
        total_stocks = len(stocks)
        
        if limit:
            stocks = stocks[:limit]
            logger.info(f"总共 {total_stocks} 只股票，本次筛选前 {limit} 只")
        else:
            logger.info(f"总共 {total_stocks} 只股票")
        
        results = []
        success_count = 0
        error_count = 0
        
        for idx, stock in enumerate(stocks, 1):
            stock_code = stock['stock_code']
            stock_name = stock['stock_name']
            
            if idx % 100 == 0:
                logger.info(f"进度: {idx}/{len(stocks)} ({idx*100//len(stocks)}%)")
            
            try:
                result = self.screen_stock(stock_code, stock_name)
                if result:
                    results.append(result)
                    logger.info(f"✓ 找到符合条件的股票: {stock_code} {stock_name} "
                              f"量比:{result['vol_ratio']:.2f} "
                              f"RSI:{result['rsi']:.1f}")
                success_count += 1
                
                # 添加小延迟，避免过于频繁的数据库查询
                if idx % 50 == 0:
                    import time
                    time.sleep(0.1)
                    
            except Exception as e:
                error_count += 1
                logger.error(f"✗ 处理股票 {stock_code} 失败: {e}")
        
        logger.info("=" * 80)
        logger.info(f"筛选完成！")
        logger.info(f"处理成功: {success_count} 只")
        logger.info(f"处理失败: {error_count} 只")
        logger.info(f"符合条件: {len(results)} 只")
        logger.info("=" * 80)
        
        # 保存结果
        if results:
            self.save_results(results)
        else:
            logger.info("没有找到符合条件的股票")
        
        return results
    
    def save_results(self, results: list):
        """保存筛选结果到CSV文件"""
        if not results:
            return
        
        df = pd.DataFrame(results)
        
        # 按量比降序排序
        df = df.sort_values('vol_ratio', ascending=False)
        
        # 生成文件名
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        csv_filename = f'macd_volume_strength_{timestamp}.csv'
        txt_filename = f'stock_codes_volume_strength_{timestamp}.txt'
        
        # 保存CSV文件
        df.to_csv(csv_filename, index=False, encoding='utf-8-sig')
        logger.info(f"结果已保存到: {csv_filename}")
        
        # 保存股票代码到txt文件
        with open(txt_filename, 'w', encoding='utf-8') as f:
            for result in results:
                f.write(f"{result['stock_code']}\n")
        logger.info(f"股票代码已保存到: {txt_filename}")
        
        # 打印前10条结果
        logger.info("\n前10条筛选结果（按量比降序）:")
        logger.info("-" * 80)
        for i, result in enumerate(results[:10], 1):
            logger.info(f"{i}. {result['stock_code']} {result['stock_name']} "
                       f"收盘:{result['close_price']} "
                       f"量比:{result['vol_ratio']:.2f} "
                       f"RSI:{result['rsi']:.1f} "
                       f"涨幅:{result['change_pct']:.2f}% "
                       f"DIF:{result['macd_dif']:.4f}")


def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='MACD成交量放大强势筛选器')
    parser.add_argument('--limit', type=int, help='限制筛选的股票数量（用于测试）')
    parser.add_argument('--stock', type=str, help='筛选指定股票代码')
    
    args = parser.parse_args()
    
    screener = MACDVolumeStrengthScreener()
    
    if args.stock:
        # 筛选单只股票
        logger.info(f"筛选单只股票: {args.stock}")
        sql = "SELECT stock_name FROM new_stock_info WHERE stock_code = %s"
        stock_info = screener.db.execute_query(sql, (args.stock,))
        if stock_info:
            stock_name = stock_info[0]['stock_name']
            result = screener.screen_stock(args.stock, stock_name)
            if result:
                logger.info(f"符合条件: {result}")
            else:
                logger.info("不符合筛选条件")
        else:
            logger.error(f"未找到股票: {args.stock}")
    else:
        # 批量筛选
        screener.run_screening(limit=args.limit)


if __name__ == '__main__':
    main()

